import os
import argparse
import csv
try:
    import pydicom
    import pandas as pd
except:
    os.system("pip3 config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple")
    os.system("pip3 install pydicom pandas openpyxl")

import pydicom
import pandas as pd



def get_dicom_metadata(dicom_file_path):
    """读取DCM文件的属性并返回字典"""
    dicom_file = pydicom.dcmread(dicom_file_path, force=True)
    attributes = {}

    for elem in dicom_file:
        if elem.VR != 'SQ':
            attributes[elem.name] = elem.value
    
    return attributes

def traverse_and_rename_folders(root_folder, csv_file=None, excel_file=None):
    """遍历根文件夹并读取每个子子文件夹中的一个DCM文件，同时重命名子文件夹并更新CSV和Excel文件"""
    mapping = []
    
    # 读取现有的CSV文件内容
    if csv_file and os.path.exists(csv_file):
        with open(csv_file, mode='r', newline='') as file:
            reader = csv.reader(file)
            existing_mappings = list(reader)
        # 去掉标题行
        existing_mappings = existing_mappings[1:]
    else:
        existing_mappings = []

    # 读取Excel文件内容
    if excel_file and os.path.exists(excel_file):
        df = pd.read_excel(excel_file)
    else:
        df = None

    for subdir, dirs, _ in os.walk(root_folder):
        for dir in dirs:
            subsubdir_path = os.path.join(subdir, dir)
            for subsubdir, _, files in os.walk(subsubdir_path):
                for file in files:
                    if file.lower().endswith('.dcm'):
                        dicom_file_path = os.path.join(subsubdir, file)
                        metadata = get_dicom_metadata(dicom_file_path)
                        if metadata:
                            original_folder_name = os.path.basename(subdir)
                            parent_folder = os.path.dirname(subdir)
                                # 手动指定字符集为 GB18030 或 GB2312
                            new_folder_name = metadata["Patient's Name"]
                            try:
                                new_folder_name = new_folder_name.decode('gb18030')
                            except AttributeError:
                                # 如果已经是字符串类型，则跳过解码
                                pass
                            new_folder_name = str(new_folder_name)
                            new_folder_path = os.path.join(parent_folder, new_folder_name)

                            # 检查原文件夹名是否已经在B列中存在
                            already_exists = any(new_folder_name == row[1] for row in existing_mappings)
                            
                            if not already_exists:
                                # 重命名文件夹
                                try:
                                    os.rename(subdir, new_folder_path)
                                except:
                                    pass

                                # 记录原文件夹名称和新文件夹名称
                                mapping.append((original_folder_name, new_folder_name))

                                # print(f"文件夹 {original_folder_name} 重命名为 {new_folder_name}")

                                # 更新Excel文件
                                if df is not None:
                                    count = 1
                                    for idx, row in df.iterrows():
                                        if row.iloc[5] == new_folder_name:
                                            df.at[idx, df.columns[0]] = f"D{count}"
                                            count += 1
                        break  # 只读取一个DCM文件
                break  # 只处理一个子子文件夹

    # 将新的映射关系写入CSV文件
    if csv_file:
        with open(csv_file, mode='a', newline='') as file:
            writer = csv.writer(file)
            if not existing_mappings:  # 如果CSV文件是新的，写入标题行
                writer.writerow(["Original Folder Name", "Patient Name"])
            writer.writerows(mapping)

    # 保存更新后的Excel文件
    if df is not None:
        df.to_excel(excel_file, index=False)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="读取DCM文件属性并重命名文件夹，同时更新CSV和Excel文件")
    parser.add_argument("root_folder", type=str, help="根文件夹路径")
    parser.add_argument("--csv_file", type=str, help="保存映射关系的CSV文件路径", default=None)
    parser.add_argument("--excel_file", type=str, help="需要更新的Excel文件路径", default=None)
    args = parser.parse_args()

    # 遍历根文件夹并读取DCM文件属性
    traverse_and_rename_folders(args.root_folder, args.csv_file, args.excel_file)
